import cv2
from flask import Flask, request, jsonify
import tempfile

from src.common_util import CommonUtil
from src.detector import Detector

import os

from src.result import response_json, response_error
from src.result_handle import process_helmet_video, detect_image_blur, detect_video_blur, predict_image
from src.ms_server_api import MsServerApi

ms_server_api = MsServerApi()
detector = Detector()
common_util = CommonUtil()

app = Flask(__name__)


@app.route('/')
def hello_world():  # put application's code here
    return 'Hello World!'


@app.route('/predict/fire', methods=['POST'])
def predict_fire():
    """火焰识别"""
    file = request.files.get("file")
    if not file:
        return jsonify({"error": "请上传图片或视频"}), 400
    # 保存到临时文件
    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as tmp:
        tmp.write(file.read())
        sourcefile_tmp_path = tmp.name  # 临时文件路径

    try:
        name, ext = os.path.splitext(file.filename)
        detected_filename = "uploads/fire_predict_file" + ext + ext

        if common_util.check_file_type(ext) == 'image':
            img_oss, video_oss, full_oss = predict_image(sourcefile_tmp_path, "fire"), None, None
        elif common_util.check_file_type(ext) == 'video':
            img_oss, video_oss, full_oss = process_helmet_video(sourcefile_tmp_path, detected_filename, "fire")
        else:
            os.remove(sourcefile_tmp_path)
            return jsonify({"error": "不支持的文件类型"}), 400
    finally:
        # 删除临时文件，避免堆积
        if os.path.exists(sourcefile_tmp_path):
            os.remove(sourcefile_tmp_path)
    print(">>>>>>>> fire detected success.")
    print({
        "img_oss": img_oss,
        "video_oss": video_oss,
        "full_oss": full_oss
    })
    return response_json({
        "img_oss": img_oss,
        "video_oss": video_oss,
        "full_oss": full_oss
    })


@app.route('/predict/head', methods=['POST'])
def predict_head():
    """头盔识别"""
    file = request.files.get("file")
    if not file:
        return jsonify({"error": "请上传图片或视频"}), 400
    # 保存到临时文件
    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as tmp:
        tmp.write(file.read())
        sourcefile_tmp_path = tmp.name  # 临时文件路径

    try:
        name, ext = os.path.splitext(file.filename)
        detected_filename = "uploads/helmet_predict_file" + ext + ext
        if common_util.check_file_type(ext) == 'image':
            img_oss, video_oss, full_oss = predict_image(sourcefile_tmp_path, "head"), None, None
        elif common_util.check_file_type(ext) == 'video':
            img_oss, video_oss, full_oss = process_helmet_video(sourcefile_tmp_path, detected_filename, "head")
        else:
            os.remove(sourcefile_tmp_path)
            return jsonify({"error": "不支持的文件类型"}), 400

    finally:
        # 删除临时文件，避免堆积
        if os.path.exists(sourcefile_tmp_path):
            os.remove(sourcefile_tmp_path)
    print(">>>>>>>> head detected success.")
    print({
        "img_oss": img_oss,
        "video_oss": video_oss,
        "full_oss": full_oss
    })
    return response_json({
        "img_oss": img_oss,
        "video_oss": video_oss,
        "full_oss": full_oss
    })


@app.route('/predict/person', methods=['POST'])
def predict_person():
    """人体识别"""
    file = request.files.get("file")
    if not file:
        return jsonify({"error": "请上传图片或视频"}), 400
    # 保存到临时文件
    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as tmp:
        tmp.write(file.read())
        sourcefile_tmp_path = tmp.name  # 临时文件路径

    try:
        name, ext = os.path.splitext(file.filename)
        detected_filename = "uploads/person_predict_file" + ext + ext
        if common_util.check_file_type(ext) == 'image':
            img_oss, video_oss, full_oss = predict_image(sourcefile_tmp_path, "person"), None, None
        elif common_util.check_file_type(ext) == 'video':
            img_oss, video_oss, full_oss = process_helmet_video(sourcefile_tmp_path, detected_filename, "person")
        else:
            os.remove(sourcefile_tmp_path)
            return jsonify({"error": "不支持的文件类型"}), 400
    finally:
        # 删除临时文件，避免堆积
        if os.path.exists(sourcefile_tmp_path):
            os.remove(sourcefile_tmp_path)

    print(">>>>>>>> person detected success.")
    print({
        "img_oss": img_oss,
        "video_oss": video_oss,
        "full_oss": full_oss
    })
    return response_json({
        "img_oss": img_oss,
        "video_oss": video_oss,
        "full_oss": full_oss
    })


@app.route('/predict/obscure', methods=['POST'])
def predict_obscure():
    """模糊/遮挡检测"""
    file = request.files.get("file")
    if not file:
        return jsonify({"error": "请上传图片或视频"}), 400

    # 保存到临时文件
    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as tmp:
        file.save(tmp.name)
        tmp_path = tmp.name

    # 判断文件是图片还是视频
    ext = os.path.splitext(file.filename)[1].lower()
    if common_util.check_file_type(ext) == 'image':
        img_oss = '无法识别图片是否遮挡'  # detect_image_blur(tmp_path)
    elif common_util.check_file_type(ext) == 'video':
        img_oss = detect_video_blur(tmp_path)
    else:
        os.remove(tmp_path)
        return response_error("不支持的文件类型")

    # 删除临时文件
    os.remove(tmp_path)

    return response_json({
        "img_oss": img_oss
    })


def variance_of_laplacian(gray):
    return cv2.Laplacian(gray, cv2.CV_64F).var()


def parse_results(results):
    """
    解析 YOLO 的预测结果为 JSON
    """
    data = []
    for r in results:
        for box in r.boxes:
            cls = int(box.cls[0])  # 类别ID
            conf = float(box.conf[0])  # 置信度
            xyxy = box.xyxy[0].tolist()  # 坐标
            data.append({
                "class": r.names[cls],
                "confidence": conf,
                "bbox": xyxy
            })
    return {"results": data}


if __name__ == '__main__':
    # app.run(host="8.156.82.84", port=5055)
    app.run(host="0.0.0.0", port=5055)
    # app.run(host="47.109.52.29", port=5000)
